

#include <kmalloc.h>
#include <types.h>
#include <linkedlist.h>
#include <debug.h>

// 4k of 8-byte mem
static size_t m8_length = M4K/8;  // length of stack
static void * m8_stack[M4K/8];  // stack of free addresses
static void ** m8_head;  // stack pointer
static byte m8_free[M4K];  // pre-allocated space

// 4k of 32-byte mem
static size_t m16_length = M8K/16;
static void * m16_stack[M8K/16];
static void ** m16_head;
static byte m16_free[M8K];

// 8k of 64-byte mem
static size_t m64_length = M8K/64;
static void * m64_stack[M8K/64];
static void ** m64_head;
static byte m64_free[M8K];

// 8k of 256-byte mem
static size_t m256_length = M8K/256;
static void * m256_stack[M8K/256];
static void ** m256_head;
static byte m256_free[M8K];

/**
 * Initialize the Kernel Memory Allocator
 */
int init_kmalloc(void)
{
	int ret = 0;
	
	int i;
	m8_head = &m8_stack[M4K/8];
	for (i = 0; i < m8_length; i++)
		m8_stack[i] = &m8_free[i*8];
	
	m16_head = m16_stack + m16_length;
	for (i = 0; i < m16_length; i++)
		m16_stack[i] = &m16_free[i*16];
	
	m64_head = m64_stack + m64_length;
	for (i = 0; i < m64_length; i++)
		m64_stack[i] = &m64_free[i*64];
	
	m256_head = m256_stack + m256_length;
	for (i = 0; i < m256_length; i++)
		m256_stack[i] = &m256_free[i*256];
	
	debugf("[KMALLOC] 8-byte stack @ 0x%08x, head @ 0x%08x, free* @ 0x%08x\n", m8_stack, m8_head, m8_free);
	
	return ret;
}

/**
 * Allocate a chunk of memory of a given size
 * 
 * size_t size - size of memory chunk requested
 * 
 * returns - pointer to the memory chunck, or NULL on error
 */
void * kmalloc(size_t size)
{
	void *ptr = NULL;
	
	if (size <= 8)
	{
		if ((m8_head - 1) >= m8_stack)
		{
			m8_head--;
			ptr = *m8_head;
			debugf("[KMALLOC] Allocated 8-bytes to 0x%08x, head = 0x%08x\n", ptr, m8_head);
			return ptr;
		}
	}
	
	if (size <= 16)
	{
		if ((m16_head - 1) >= m16_stack)
		{
			m16_head--;
			ptr = *m16_head;
			debugf("[KMALLOC] Allocated 32-bytes to 0x%08x, head = 0x%08x\n", ptr, m16_head);
			return ptr;
		}
	}
	
	if (size <= 64)
	{
		if ((m64_head - 1) >= m64_stack)
		{
			m64_head--;
			ptr = *m64_head;
			debugf("[KMALLOC] Allocated 64-bytes to 0x%08x, head = 0x%08x\n", ptr, m64_head);
			return ptr;
		}
	}
	
	if (size <= 256)
	{
		if ((m256_head - 1) >= m256_stack)
		{
			m256_head--;
			ptr = *m256_head;
			debugf("[KMALLOC] Allocated 256-bytes to 0x%08x, head = 0x%08x\n", ptr, m256_head);
			return ptr;
		}
	}
	
	debugf("[KMALLOC] Unhandled size: %d or all stacks are full!\n", size);
	
	return ptr;
}

/**
 * Frees a chunk of memory previously allocated with kmalloc
 * 
 * void *ptr - pointer to the memory chunk
 * 
 * TODO: may want to check if the *ptr is already on the free stack it catches if you try
 * to free memory more than once... slow?
 */
void kfree(void *ptr)
{
	debugf("[KFREE] ptr: %x, m8_free: %x, m8_free+4k: %x\n", ptr, (void*)m8_free, (void*)m8_free+M4K);
	if ((ptr >= (void*)m8_free) && (ptr < ((void*)m8_free + M4K)))
	{
		debugf("[KFREE] Freed 8-bytes\n");
		*m8_head = ptr;
		m8_head += sizeof(void*);
	}
	else if ((ptr >= (void*)m16_free) && (ptr < ((void*)m16_free + M8K)))
	{
		debugf("[KFREE] Freed 16-bytes\n");
		*m16_head = ptr;
		m16_head += sizeof(void*);
	}
	else if ((ptr >= (void*)m64_free) && (ptr < ((void*)m64_free + M8K)))
	{
		debugf("[KFREE] Freed 64-bytes\n");
		*m64_head = ptr;
		m64_head += sizeof(void*);
	}
	else if ((ptr >= (void*)m256_free) && (ptr < ((void*)m256_free + M8K)))
	{
		debugf("[KFREE] Freed 256-bytes\n");
		*m256_head = ptr;
		m256_head += sizeof(void*);
	}
	else
	{
		debugf("[KFREE] Unknown ptr 0x%08x\n", ptr);
	}
}
